#!/usr/bin/env python3
from turtle import left
import numpy as np
from typing import Dict, Any
import sys
import os
from pathlib import Path
# Fix for Python 3.10 compatibility - ensure path is string
robottwin_path = str(Path(__file__).parent.parent / 'RoboTwin')
if robottwin_path not in sys.path:
    sys.path.append(robottwin_path)
from envs.utils.vis_pcd import PointArrow, catpcd
from scalePcd import scalePcd
from keyPcd import keypose_item
from scipy.spatial.transform import Rotation

def trajPcd(unifiedData: Dict[str, Any], param: str = 's', *, surpress_print: bool = False) -> Dict[str, Any]:
    """
    Convert trajectory data (state/action) to point cloud arrows and inject into existing pointclouds.
    about left and right: the action / state 's 0:7 is the left on first view, and keypose's left is also left on first view
    but the semantic part's left is third view's left, so here the current_bias l and r are second view's and when adding to state/action, it needs to be crossed.
    """
    if 'key_poses' in unifiedData and unifiedData.get("key_poses",None) is not None:
        ee_state_data = unifiedData.get("ee_state", None)
        assert ee_state_data is not None, "No ee_state data available"
        left_state=np.array(ee_state_data[:3])
        right_state=np.array(ee_state_data[8:11])
        state_data_to_check = ee_state_data
        assert left_state.shape == (3,) and right_state.shape == (3,), f"Expected left_state and right_state to be (3,), got {left_state.shape} and {right_state.shape}, if zero shaped, perhaps not added as keys to load"
        keypose=np.array(unifiedData.get("key_poses",None))
        left_pos=np.array(keypose_item(keypose)["left_gripper"]["pos"])
        right_pos=np.array(keypose_item(keypose)["right_gripper"]["pos"])
        # print("left_pos",left_pos.shape,"left_state",left_state.shape,"right_pos",right_pos.shape,"right_state",right_state.shape,"keypose",keypose.shape)
        # print("left_pos",left_pos,"left_state",left_state,"right_pos",right_pos,"right_state",right_state)
        # if not np.allclose(left_pos,left_state,atol=0.3) or not np.allclose(right_pos,right_state,atol=0.3):
        #     print(list(np.round(state_data_to_check,3)),"seems not an ee pose, but rather qpos")
        #     return unifiedData
    if param in ["s","a"]:
        return add_js_ja(unifiedData, param, surpress_print)
    elif param == 'o':
        return add_obj_pose(unifiedData,param,surpress_print)
    elif param == "e":
        return add_mid_end(unifiedData, param, surpress_print)
    elif param.startswith("m"):
        return move_as(unifiedData, param[1:], surpress_print,"")
    elif param.startswith("i"):
        return move_as(unifiedData, param[1:], surpress_print, "inv")
    print(f"    ❌ TRAJPCD({param}): Invalid parameter, got {param}")
    return unifiedData 
def move_as_7D(state_7D, local_vec, flag):
    """Move 7D pose by local vector in local frame, return modified global pose"""
    pos, quat = state_7D[:3], state_7D[3:]
    local_vec=np.array(local_vec)
    if flag == "inv":
        rotation = Rotation.from_quat(quat).inv()
    else:
        rotation = Rotation.from_quat(quat)
    rotated_vec = rotation.apply(local_vec)
    return np.concatenate([pos + rotated_vec, quat])
def move_as_basic(unifiedData: Dict[str, Any],bias_vec_l,bias_vec_r, flag) -> Dict[str, Any]:
    print("unifiedData['ee_state']",list(np.round(unifiedData['ee_state'],3)),"bias_vec_l",list(np.round(bias_vec_l,3)),"bias_vec_r",list(np.round(bias_vec_r,3)))
    assert bias_vec_l is not None and bias_vec_r is not None, f"No stick_bias_vec data available for move_as, please add eng* into pipeline first"
    unifiedData['ee_state'][:7] = move_as_7D(unifiedData['ee_state'][:7], bias_vec_l, flag)
    unifiedData['ee_state'][8:15] = move_as_7D(unifiedData['ee_state'][8:15], bias_vec_r, flag)
    print("moved ee_state",flag,list(np.round(unifiedData['ee_state'][:3],3)),list(np.round(unifiedData['ee_state'][8:11],3)))
    unifiedData['ee_action'][:7] = move_as_7D(unifiedData['ee_action'][:7], bias_vec_l, flag)
    unifiedData['ee_action'][8:15] = move_as_7D(unifiedData['ee_action'][8:15], bias_vec_r, flag)
    print("moved ee_action",flag,list(np.round(unifiedData['ee_action'][:3],3)),list(np.round(unifiedData['ee_action'][8:11],3)))
def move_as(unifiedData: Dict[str, Any],param,surpress_print,flag) -> Dict[str, Any]:
    """Move both arms' states by local vectors"""
    if param== 's':
        move_as_basic(unifiedData, unifiedData['stick_bias_vec'], unifiedData['stick_bias_vec'], flag)
    elif param== 'e':
        bias_vec_l, bias_vec_r = get_js_je_diff(unifiedData)
        move_as_basic(unifiedData, bias_vec_r, bias_vec_l, flag)
    return unifiedData
def get_js_je_diff(unifiedData: Dict[str, Any]) -> Dict[str, Any]:
    lend_midpose, rend_midpose = get_mid_end_pose(unifiedData)
    left_ee,right_ee = unifiedData['ee_state'][:3],unifiedData['ee_state'][8:11]
    return np.concatenate([rend_midpose[0:3] - left_ee,rend_midpose[3:]]) , np.concatenate([lend_midpose[0:3] - right_ee,lend_midpose[3:]])
    return rend_midpose[0:3] - left_ee , lend_midpose[0:3] - right_ee
def get_mid_end_pose(unifiedData: Dict[str, Any]) -> Dict[str, Any]:
    assert "key_poses" in unifiedData and unifiedData["key_poses"] is not None, f"No key_poses data available for trajPcd"
    keypose_data = unifiedData['key_poses']
    llend=keypose_item(keypose_data)["left_gripper"]["left_finger_end"]
    lrend=keypose_item(keypose_data)["left_gripper"]["right_finger_end"]
    rlend=keypose_item(keypose_data)["right_gripper"]["left_finger_end"]
    rrend=keypose_item(keypose_data)["right_gripper"]["right_finger_end"]
    lend_mid_xyz=((np.array(lrend)+np.array(llend))/2)
    lend_mid_rpy=keypose_item(keypose_data)["left_gripper"]["rpy"]
    rend_mid_xyz=((np.array(rlend)+np.array(rrend))/2)
    rend_mid_rpy=keypose_item(keypose_data)["right_gripper"]["rpy"]
    lend_midpose=np.concatenate([lend_mid_xyz,lend_mid_rpy])
    rend_midpose=np.concatenate([rend_mid_xyz,rend_mid_rpy])
    return lend_midpose, rend_midpose
def add_mid_end(unifiedData: Dict[str, Any], param: str = 'e',  surpress_print: bool = False) -> Dict[str, Any]:
    lend_midpose, rend_midpose = get_mid_end_pose(unifiedData)
    mid_end_pcd_1=PointArrow(lend_midpose, [0, 255, 0], 100)  # green
    mid_end_pcd_2=PointArrow(rend_midpose, [0, 0, 255], 100)  # blue
    mid_end_pcd=catpcd(mid_end_pcd_1, mid_end_pcd_2)
    return add_arrow_u(unifiedData, mid_end_pcd, param, surpress_print)
def add_js_ja(unifiedData: Dict[str, Any], param: str ,surpress_print: bool = False) -> Dict[str, Any]:
    data_key = {'s': 'ee_state', 'a': 'ee_action'}[param]
    assert data_key in unifiedData and unifiedData[data_key] is not None, f"No {data_key} data available for trajPcd"
    trajectory_data = np.array(unifiedData[data_key])
    assert len(trajectory_data.shape) == 1, f"Expected 1D trajectory data, got shape {trajectory_data.shape}"
    assert trajectory_data.shape[0] in [16,8] ,f"Expected 16D trajectory data, got shape {trajectory_data.shape}"
    pose_data_1 = trajectory_data[:7]  # Take first 7 elements as pose
    pose_data_2 = trajectory_data[8:15]  # Take first 7 elements as pose
    traj_pcd_1 = PointArrow(pose_data_1, [255, 0, 255] if param == 's' else [0, 255, 255], 100)  # Magenta for state, Cyan for action on the 0:7 side of 16 D
    traj_pcd_2 = PointArrow(pose_data_2, [127, 0, 127] if param == 's' else [0, 127, 127], 100)  # Dark red = for state,Dark blue for action on the 8:15 side of 16 D
    traj_pcd = catpcd(traj_pcd_1, traj_pcd_2)
    return add_arrow_u(unifiedData, traj_pcd, param, surpress_print)
def add_obj_pose(unifiedData: Dict[str, Any], param: str, surpress_print: bool = False) -> Dict[str, Any]:
    assert "foundationPoses" in unifiedData and unifiedData["foundationPoses"] is not None, f"No foundationPoses data available for trajPcd"
    for i,obj_pose in enumerate(unifiedData['foundationPoses']):
        obj_pcd = PointArrow(obj_pose, [255*(i==0),255*(i==1),255*(i==2)], 100)
        print("obj_pose",i,obj_pose,"obj_pcd",obj_pcd.shape,"color",obj_pcd[0,3:6])
        unifiedData = add_arrow_u(unifiedData, obj_pcd, param, surpress_print)
    return unifiedData
def add_arrow_u(unifiedData: Dict[str, Any], traj_pcd, param: str,surpress_print: bool = False):
    traj_pcd = scalePcd(traj_pcd, tag=f"traj_{param}")
    updated_fields = {}
    for pcd_key in unifiedData.keys():
        if "pointcloud" in pcd_key and unifiedData[pcd_key] is not None:
            existing_pcd = unifiedData[pcd_key]
            assert len(existing_pcd.shape) == 2 and existing_pcd.shape[1] == 6, f"Skipping {pcd_key}: Expected shape (Np, 6), got {existing_pcd.shape}"
            original_count = existing_pcd.shape[0]
            existing_pcd = scalePcd(existing_pcd, tag="existing_pcd")
            traj_pcd_to_use = scalePcd(traj_pcd, tag="traj_pcd_to_use")
            combined_pcd = catpcd(existing_pcd, traj_pcd_to_use)
            unifiedData[pcd_key] = combined_pcd
            updated_fields[pcd_key] = (original_count, combined_pcd.shape[0])
    
    for key, (original_count, final_count) in updated_fields.items():
        if not surpress_print:
            print(f"    🎯 TRAJPCD({param},{key}): {original_count} → {final_count}")
    return unifiedData